import tkinter as tk
import numpy as np
from tkinter import messagebox, Button, Canvas, Label, Frame
from PIL import Image, ImageDraw
import os
import pickle

try:
    import tensorflow as tf
    from tensorflow.keras.models import load_model
    TF_AVAILABLE = True
except ImportError:
    TF_AVAILABLE = False

class DigitRecognitionApp:
    def __init__(self, root):
        self.root = root
        self.root.title("手写数字识别")
        self.root.geometry("800x600")
        self.root.resizable(False, False)
        
        # 设置变量
        self.canvas_width = 280
        self.canvas_height = 280
        self.brush_size = 15
        self.model = None
        
        # 创建界面
        self.create_widgets()
        
        # 加载模型
        self.load_recognition_model()
    
    def create_widgets(self):
        # 创建主框架
        main_frame = Frame(self.root)
        main_frame.pack(fill="both", expand=True, padx=20, pady=20)
        
        # 左侧框架 - 绘图区域
        left_frame = Frame(main_frame, bd=2, relief="ridge", padx=10, pady=10)
        left_frame.pack(side="left", fill="both", expand=True)
        
        # 右侧框架 - 控制和结果区域
        right_frame = Frame(main_frame, bd=2, relief="ridge", padx=10, pady=10)
        right_frame.pack(side="right", fill="both", expand=True)
        
        # 创建画布
        self.canvas = Canvas(left_frame, width=self.canvas_width, height=self.canvas_height, 
                            bg="black", cursor="pencil")
        self.canvas.pack(pady=20)
        
        # 绑定鼠标事件
        self.canvas.bind("<B1-Motion>", self.paint)
        
        # 创建控制按钮
        btn_frame = Frame(left_frame)
        btn_frame.pack(pady=10)
        
        clear_btn = Button(btn_frame, text="清除", width=10, command=self.clear_canvas)
        clear_btn.grid(row=0, column=0, padx=5)
        
        recognize_btn = Button(btn_frame, text="识别", width=10, command=self.recognize_digit)
        recognize_btn.grid(row=0, column=1, padx=5)
        
        # 右侧 - 结果显示区域
        result_label = Label(right_frame, text="识别结果", font=("Arial", 16))
        result_label.pack(pady=10)
        
        self.result_display = Label(right_frame, text="-", font=("Arial", 72, "bold"))
        self.result_display.pack(pady=20)
        
        # 置信度显示
        confidence_frame = Frame(right_frame)
        confidence_frame.pack(pady=10, fill="both", expand=True)
        
        Label(confidence_frame, text="各数字的置信度", font=("Arial", 12)).pack()
        
        self.confidence_labels = []
        for i in range(10):
            row = i // 5
            col = i % 5
            frame = Frame(confidence_frame)
            frame.pack(side="left" if col < 5 else "right", padx=5, pady=5)
            
            Label(frame, text=f"数字 {i}:", font=("Arial", 10)).pack()
            conf_label = Label(frame, text="0%", font=("Arial", 10))
            conf_label.pack()
            self.confidence_labels.append(conf_label)
        
        # 状态显示
        self.status_label = Label(self.root, text="就绪", bd=1, relief="sunken", anchor="w")
        self.status_label.pack(side="bottom", fill="x")
        
        # 创建图像对象用于保存绘图
        self.image = Image.new("L", (self.canvas_width, self.canvas_height), color=0)
        self.draw = ImageDraw.Draw(self.image)
    
    def paint(self, event):
        x1, y1 = (event.x - self.brush_size), (event.y - self.brush_size)
        x2, y2 = (event.x + self.brush_size), (event.y + self.brush_size)
        
        # 在Canvas上绘制
        self.canvas.create_oval(x1, y1, x2, y2, fill="white", outline="white")
        
        # 同时在PIL图像上绘制
        self.draw.ellipse([x1, y1, x2, y2], fill=255)
    
    def clear_canvas(self):
        self.canvas.delete("all")
        self.image = Image.new("L", (self.canvas_width, self.canvas_height), color=0)
        self.draw = ImageDraw.Draw(self.image)
        self.result_display.config(text="-")
        for label in self.confidence_labels:
            label.config(text="0%")
        self.status_label.config(text="就绪")
    
    def load_recognition_model(self):
        try:
            if TF_AVAILABLE:
                # 检查模型文件是否存在
                if os.path.exists("digit_model.h5"):
                    self.model = load_model("digit_model.h5")
                    self.status_label.config(text="模型加载成功")
                else:
                    self.status_label.config(text="模型文件不存在，将在首次使用时训练")
            else:
                self.status_label.config(text="TensorFlow未安装，将使用简单模型进行识别")
                # 尝试加载简单模型
                if os.path.exists("simple_model.pkl"):
                    with open("simple_model.pkl", "rb") as f:
                        self.model = pickle.load(f)
        except Exception as e:
            messagebox.showerror("错误", f"加载模型时出错: {str(e)}")
            self.status_label.config(text="模型加载失败")
    
    def preprocess_image(self):
        try:
            # 调整图像大小为28x28（MNIST格式）
            img = self.image.copy()
            
            # 兼容不同版本的PIL库
            try:
                # 新版本PIL
                img = img.resize((28, 28), Image.Resampling.LANCZOS)
            except AttributeError:
                # 旧版本PIL
                img = img.resize((28, 28), Image.LANCZOS)
            except Exception:
                # 最后的备选方案
                img = img.resize((28, 28))
            
            # 转换为numpy数组并归一化
            img_array = np.array(img)
            img_array = img_array / 255.0
            
            if TF_AVAILABLE and self.model:
                # 为TensorFlow模型准备输入
                return img_array.reshape(1, 28, 28, 1)
            else:
                # 为简单模型准备输入
                return img_array.reshape(1, 784)
        except Exception as e:
            messagebox.showerror("错误", f"图像预处理失败: {str(e)}")
            self.status_label.config(text="图像预处理失败")
            raise
    
    def recognize_digit(self):
        if self.model is None:
            self.status_label.config(text="模型未加载，尝试训练新模型...")
            self.root.update()
            if TF_AVAILABLE:
                self.train_model()
            else:
                self.train_simple_model()
            
            # 再次检查模型是否成功加载
            if self.model is None:
                messagebox.showerror("错误", "模型训练失败，无法进行识别")
                self.status_label.config(text="模型训练失败")
                return
        
        try:
            # 预处理图像
            self.status_label.config(text="正在预处理图像...")
            self.root.update()
            processed_img = self.preprocess_image()
            
            # 检查处理后的图像是否有效
            if processed_img is None:
                return
                
            # 进行预测
            self.status_label.config(text="正在进行预测...")
            self.root.update()
            
            if TF_AVAILABLE and isinstance(self.model, tf.keras.Model):
                predictions = self.model.predict(processed_img)[0]
            else:
                # 简单模型预测
                try:
                    predictions = self.model.predict_proba(processed_img)[0]
                except AttributeError:
                    # 如果模型没有predict_proba方法，尝试使用predict
                    raw_predictions = self.model.predict(processed_img)
                    # 确保predictions是概率分布
                    if hasattr(self.model, "classes_"):
                        predictions = np.zeros(len(self.model.classes_))
                        predictions[raw_predictions[0]] = 1.0
                    else:
                        predictions = raw_predictions[0]
            
            # 获取预测结果
            predicted_digit = np.argmax(predictions)
            confidence = predictions[predicted_digit] * 100
            
            # 更新UI
            self.result_display.config(text=str(predicted_digit))
            self.status_label.config(text=f"识别完成，置信度: {confidence:.2f}%")
            
            # 更新置信度标签
            for i, label in enumerate(self.confidence_labels):
                if i < len(predictions):
                    label.config(text=f"{predictions[i]*100:.1f}%")
                else:
                    label.config(text="N/A")
                
        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            messagebox.showerror("错误", f"识别过程中出错: {str(e)}\n\n详细信息:\n{error_details}")
            self.status_label.config(text="识别失败")
    
    def train_model(self):
        try:
            self.status_label.config(text="正在训练模型...这可能需要几分钟时间")
            self.root.update()
            
            # 导入必要的库
            import os
            import urllib.request
            import gzip
            import struct
            import array
            import numpy as np
            
            # 设置数据目录为当前工作目录下的data文件夹
            data_home = os.path.join(os.getcwd(), "data")
            os.makedirs(data_home, exist_ok=True)
            
            # 定义MNIST数据集文件路径
            train_images_path = os.path.join(data_home, 'train-images-idx3-ubyte.gz')
            train_labels_path = os.path.join(data_home, 'train-labels-idx1-ubyte.gz')
            test_images_path = os.path.join(data_home, 't10k-images-idx3-ubyte.gz')
            test_labels_path = os.path.join(data_home, 't10k-labels-idx1-ubyte.gz')
            
            # 下载MNIST数据集
            try:
                self.status_label.config(text="正在下载MNIST数据集...")
                self.root.update()
                
                # 定义多个可用的MNIST数据集镜像站点
                mnist_mirrors = [
                    'http://yann.lecun.com/exdb/mnist/',
                    'https://storage.googleapis.com/cvdf-datasets/mnist/',
                    'https://ossci-datasets.s3.amazonaws.com/mnist/'
                ]
                
                # 下载函数，带有重试功能
                def download_with_retry(url, filepath, max_retries=3):
                    for retry in range(max_retries):
                        try:
                            self.status_label.config(text=f"正在下载: {url} (尝试 {retry+1}/{max_retries})")
                            self.root.update()
                            urllib.request.urlretrieve(url, filepath)
                            return True
                        except Exception as e:
                            if retry == max_retries - 1:
                                self.status_label.config(text=f"下载失败: {url}")
                                self.root.update()
                                print(f"下载失败: {url}, 错误: {str(e)}")
                    return False
                
                # 尝试从不同镜像下载文件
                def try_mirrors_download(filename, filepath):
                    for mirror in mnist_mirrors:
                        url = mirror + filename
                        if download_with_retry(url, filepath):
                            return True
                    return False
                
                # 检查文件是否存在，如果不存在则尝试下载
                files_to_download = [
                    (train_images_path, 'train-images-idx3-ubyte.gz'),
                    (train_labels_path, 'train-labels-idx1-ubyte.gz'),
                    (test_images_path, 't10k-images-idx3-ubyte.gz'),
                    (test_labels_path, 't10k-labels-idx1-ubyte.gz')
                ]
                
                download_success = True
                for filepath, filename in files_to_download:
                    if not os.path.exists(filepath):
                        if not try_mirrors_download(filename, filepath):
                            download_success = False
                            break
                
                if not download_success:
                    raise Exception("无法下载MNIST数据集，请检查网络连接或手动下载数据集到data目录")
                
                # 读取训练图像
                with gzip.open(train_images_path, 'rb') as f:
                    magic, size, rows, cols = struct.unpack(">IIII", f.read(16))
                    train_images = array.array("B", f.read())
                    x_train = np.array(train_images).reshape(size, rows, cols)
                
                # 读取训练标签
                with gzip.open(train_labels_path, 'rb') as f:
                    magic, size = struct.unpack(">II", f.read(8))
                    train_labels = array.array("B", f.read())
                    y_train = np.array(train_labels)
                
                # 读取测试图像
                with gzip.open(test_images_path, 'rb') as f:
                    magic, size, rows, cols = struct.unpack(">IIII", f.read(16))
                    test_images = array.array("B", f.read())
                    x_test = np.array(test_images).reshape(size, rows, cols)
                
                # 读取测试标签
                with gzip.open(test_labels_path, 'rb') as f:
                    magic, size = struct.unpack(">II", f.read(8))
                    test_labels = array.array("B", f.read())
                    y_test = np.array(test_labels)
                
                # 归一化数据
                x_train, x_test = x_train / 255.0, x_test / 255.0
                
                # 调整数据形状
                x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
                x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)
                
                self.status_label.config(text="MNIST数据集加载成功")
                self.root.update()
                
                # 创建模型
                self.status_label.config(text="正在创建模型...")
                self.root.update()
                model = tf.keras.models.Sequential([
                    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
                    tf.keras.layers.MaxPooling2D((2, 2)),
                    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
                    tf.keras.layers.MaxPooling2D((2, 2)),
                    tf.keras.layers.Flatten(),
                    tf.keras.layers.Dense(128, activation='relu'),
                    tf.keras.layers.Dropout(0.2),
                    tf.keras.layers.Dense(10, activation='softmax')
                ])
                
                # 编译模型
                model.compile(optimizer='adam',
                            loss='sparse_categorical_crossentropy',
                            metrics=['accuracy'])
                
                # 训练模型
                self.status_label.config(text="正在训练模型...这可能需要几分钟时间")
                self.root.update()
                model.fit(x_train, y_train, epochs=5, batch_size=64, validation_data=(x_test, y_test))
                
                # 保存模型
                self.status_label.config(text="正在保存模型...")
                self.root.update()
                model.save("digit_model.h5")
                self.model = model
                
                # 不再需要清理临时目录，因为我们使用本地data目录
                
                self.status_label.config(text="模型训练完成并保存")
            except Exception as dataset_error:
                messagebox.showerror("错误", f"加载MNIST数据集时出错: {str(dataset_error)}")
                self.status_label.config(text="MNIST数据集加载失败")
                raise
        except Exception as e:
            messagebox.showerror("错误", f"训练模型时出错: {str(e)}")
            self.status_label.config(text="模型训练失败")
    
    def train_simple_model(self):
        try:
            from sklearn.datasets import fetch_openml
            from sklearn.ensemble import RandomForestClassifier
            from sklearn.neural_network import MLPClassifier
            import os
            import urllib.request
            import gzip
            import struct
            import array
            
            self.status_label.config(text="正在训练简单模型...这可能需要几分钟时间")
            self.root.update()
            
            # 设置数据目录为当前工作目录下的data文件夹
            data_home = os.path.join(os.getcwd(), "data")
            os.makedirs(data_home, exist_ok=True)
            
            # 尝试使用本地下载的MNIST数据集
            try:
                # 定义MNIST数据集文件路径
                train_images_path = os.path.join(data_home, 'train-images-idx3-ubyte.gz')
                train_labels_path = os.path.join(data_home, 'train-labels-idx1-ubyte.gz')
                
                # 定义多个可用的MNIST数据集镜像站点
                mnist_mirrors = [
                    'http://yann.lecun.com/exdb/mnist/',
                    'https://storage.googleapis.com/cvdf-datasets/mnist/',
                    'https://ossci-datasets.s3.amazonaws.com/mnist/'
                ]
                
                # 下载函数，带有重试功能
                def download_with_retry(url, filepath, max_retries=3):
                    for retry in range(max_retries):
                        try:
                            self.status_label.config(text=f"正在下载: {url} (尝试 {retry+1}/{max_retries})")
                            self.root.update()
                            urllib.request.urlretrieve(url, filepath)
                            return True
                        except Exception as e:
                            if retry == max_retries - 1:
                                self.status_label.config(text=f"下载失败: {url}")
                                self.root.update()
                                print(f"下载失败: {url}, 错误: {str(e)}")
                    return False
                
                # 尝试从不同镜像下载文件
                def try_mirrors_download(filename, filepath):
                    for mirror in mnist_mirrors:
                        url = mirror + filename
                        if download_with_retry(url, filepath):
                            return True
                    return False
                
                # 如果文件不存在，则尝试下载
                if not (os.path.exists(train_images_path) and os.path.exists(train_labels_path)):
                    self.status_label.config(text="正在下载MNIST数据集...")
                    self.root.update()
                    
                    # 尝试下载训练图像和标签
                    download_success = True
                    if not os.path.exists(train_images_path):
                        if not try_mirrors_download('train-images-idx3-ubyte.gz', train_images_path):
                            download_success = False
                    
                    if download_success and not os.path.exists(train_labels_path):
                        if not try_mirrors_download('train-labels-idx1-ubyte.gz', train_labels_path):
                            download_success = False
                    
                    if not download_success:
                        raise Exception("无法下载MNIST数据集，请检查网络连接或手动下载数据集到data目录")
                
                # 读取训练图像
                with gzip.open(train_images_path, 'rb') as f:
                    magic, size, rows, cols = struct.unpack(">IIII", f.read(16))
                    train_images = array.array("B", f.read())
                    train_images = np.array(train_images).reshape(size, rows * cols)
                
                # 读取训练标签
                with gzip.open(train_labels_path, 'rb') as f:
                    magic, size = struct.unpack(">II", f.read(8))
                    train_labels = array.array("B", f.read())
                    train_labels = np.array(train_labels)
                
                # 准备数据
                X = train_images / 255.0
                y = train_labels
                
                self.status_label.config(text="MNIST数据集加载成功")
                self.root.update()
                
            except Exception as data_error:
                # 如果本地加载失败，尝试使用fetch_openml作为备选方案
                self.status_label.config(text=f"本地数据加载失败，尝试在线获取...")
                self.root.update()
                X, y = fetch_openml('mnist_784', version=1, return_X_y=True, parser='auto', data_home=data_home, cache=True)
                X = X / 255.0
            
            # 使用部分数据训练简单模型
            sample_size = 10000  # 使用较小的样本量以加快训练速度
            X_sample = X[:sample_size]
            y_sample = y[:sample_size]
            
            # 创建并训练神经网络分类器（比随机森林更准确）
            model = MLPClassifier(hidden_layer_sizes=(100,), max_iter=10, alpha=1e-4,
                                solver='sgd', verbose=0, tol=1e-4, random_state=1,
                                learning_rate_init=.1)
            model.fit(X_sample, y_sample)
            
            # 保存模型
            with open("simple_model.pkl", "wb") as f:
                pickle.dump(model, f)
            
            self.model = model
            self.status_label.config(text="简单模型训练完成并保存")
        except Exception as e:
            messagebox.showerror("错误", f"训练简单模型时出错: {str(e)}")
            self.status_label.config(text="简单模型训练失败")

if __name__ == "__main__":
    root = tk.Tk()
    app = DigitRecognitionApp(root)
    root.mainloop()